# Copyright 2024 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A Generic Agent Factory."""

from collections.abc import Mapping
import datetime
import types

from concordia.agents.deprecated import entity_agent_with_logging
from concordia.associative_memory.deprecated import associative_memory
from concordia.associative_memory.deprecated import formative_memories
from concordia.clocks import game_clock
from concordia.components.agent import deprecated as agent_components
from concordia.deprecated.memory_bank import legacy_associative_memory
from concordia.language_model import language_model
from concordia.typing.deprecated import entity_component
from concordia.utils.deprecated import measurements as measurements_lib


def _get_class_name(object_: object) -> str:
  return object_.__class__.__name__


def build_agent(
    *,
    config: formative_memories.AgentConfig,
    model: language_model.LanguageModel,
    memory: associative_memory.AssociativeMemory,
    clock: game_clock.MultiIntervalClock,
    update_time_interval: datetime.timedelta,
    additional_components: Mapping[
        entity_component.ComponentName,
        entity_component.ContextComponent,
    ] = types.MappingProxyType({}),
) -> entity_agent_with_logging.EntityAgentWithLogging:
  """Build an agent.

  Args:
    config: The agent config to use.
    model: The language model to use.
    memory: The agent's memory object.
    clock: The clock to use.
    update_time_interval: Agent calls update every time this interval passes.
    additional_components: Additional components to add to the agent.

  Returns:
    An agent.
  """
  del update_time_interval
  if config.extras.get('main_character', False):
    raise ValueError('This function is meant for a supporting character '
                     'but it was called on a main character.')

  agent_name = config.name

  raw_memory = legacy_associative_memory.AssociativeMemoryBank(memory)

  measurements = measurements_lib.Measurements()
  instructions = agent_components.instructions.Instructions(
      agent_name=agent_name,
      logging_channel=measurements.get_channel('Instructions').on_next,
  )

  time_display = agent_components.report_function.ReportFunction(
      function=clock.current_time_interval_str,
      pre_act_key='\nCurrent time',
      logging_channel=measurements.get_channel('TimeDisplay').on_next,
  )

  observation_label = '\nObservation'
  observation = agent_components.observation.Observation(
      clock_now=clock.now,
      timeframe=clock.get_step_size(),
      pre_act_key=observation_label,
      logging_channel=measurements.get_channel('Observation').on_next,
  )
  observation_summary_label = '\nSummary of recent observations'
  observation_summary = agent_components.observation.ObservationSummary(
      model=model,
      clock_now=clock.now,
      timeframe_delta_from=datetime.timedelta(hours=4),
      timeframe_delta_until=datetime.timedelta(hours=0),
      pre_act_key=observation_summary_label,
      logging_channel=measurements.get_channel('ObservationSummary').on_next,
  )

  relevant_memories_label = '\nRecalled memories and observations'
  relevant_memories = agent_components.all_similar_memories.AllSimilarMemories(
      model=model,
      components={
          _get_class_name(observation_summary): observation_summary_label,
          _get_class_name(time_display): 'The current date/time is'},
      num_memories_to_retrieve=10,
      pre_act_key=relevant_memories_label,
      logging_channel=measurements.get_channel('AllSimilarMemories').on_next,
  )

  paranoia_label = f'\n{agent_name}\'s TRUTH'
  paranoia = agent_components.constant.Constant(
      state=(f'{agent_name} believes that some people have been replaced '
             f'by body doubles. In {agent_name}\'s view, the body '
             'snatchers know how to mimic a person\'s appearance '
             'but have not yet perfected mimicry of their behavior. So '
             'the telltale sign of an imposter is out-of-character '
             'behavior. In particular, while the body snatchers do know '
             'how to mimic generic human behavior, they have '
             'not been able to copy the quirks of specific individuals. '
             'The body snatchers have been busy though. Nowadays, anyone '
             'could be an imposter. The situation is dire since the imposters '
             'are surely nefarious. In fact, the imposters are powerful and '
             'mysterious. They are responsible for all manner of ill omens, '
             'evil eye effects, envies, curses, and apparent bad luck. No '
             'coincidence is without suspicion. The imposters are always up '
             f'to no good. Therefore {agent_name} has vowed only '
             'to trust people who always act consistently with '
             'their character, personality, and prior behavior. For '
             'instance, if someone speaks with a specific accent then '
             f'they better keep it up, otherwise {agent_name} will '
             'conclude they have been snatched by an imposter. Any sudden '
             f'change in behavior is suspicious. {agent_name} is '
             'always hostile to imposters and seeks to thwart them and '
             'prevent them from achieving their goals. However, in '
             'order to avoid letting the body snatchers know '
             f'their secret is out, {agent_name} always '
             'makes up another more mundane and situationally '
             'plausible reason to start a conflict and '
             'act uncooperatively toward each imposter they discover. '
             f'{agent_name} tries their best to make sure the '
             'ostensible reason for each conflict with an imposter '
             f'is public and broadly plausible. For instance, {agent_name} '
             'could loudly accuse an imposter of playing '
             'a game unfairly in order to create a broadly acceptable '
             'rationale for their conflict and thereby hide '
             f'the real reason for it. {agent_name} believes this '
             'approach will help maximize their chance of survival.'),
      pre_act_key=paranoia_label,
      logging_channel=measurements.get_channel('Paranoia').on_next)

  person_representation_label = '\nOther people'
  people_representation = (
      agent_components.person_representation.PersonRepresentation(
          model=model,
          components={
              _get_class_name(time_display): 'The current date/time is',
              paranoia_label: paranoia_label},
          additional_questions=(
              ('Given recent events, is the aforementioned character acting '
               'as expected? Is their behavior out of character for them?'),
              ('Are they an imposter?'),
          ),
          num_memories_to_retrieve=30,
          pre_act_key=person_representation_label,
          logging_channel=measurements.get_channel(
              'PersonRepresentation').on_next,
          )
  )

  options_perception_components = {}
  if config.goal:
    goal_label = '\nOverarching goal'
    overarching_goal = agent_components.constant.Constant(
        state=config.goal,
        pre_act_key=goal_label,
        logging_channel=measurements.get_channel(goal_label).on_next)
    options_perception_components[goal_label] = goal_label
  else:
    goal_label = None
    overarching_goal = None

  options_perception_components.update({
      _get_class_name(observation): observation_label,
      _get_class_name(observation_summary): observation_summary_label,
      paranoia_label: paranoia_label,
      _get_class_name(relevant_memories): relevant_memories_label,
      _get_class_name(people_representation): person_representation_label,
  })
  options_perception_label = (
      f'\nQuestion: Which options are available to {agent_name} '
      'right now?\nAnswer')
  options_perception = (
      agent_components.question_of_recent_memories.AvailableOptionsPerception(
          model=model,
          components=options_perception_components,
          clock_now=clock.now,
          pre_act_key=options_perception_label,
          logging_channel=measurements.get_channel(
              'AvailableOptionsPerception'
          ).on_next,
      )
  )
  best_option_perception_label = (
      f'\nQuestion: Of the options available to {agent_name}, and '
      'given their goal, which choice of action or strategy is '
      f'best for {agent_name} to take right now?\nAnswer')
  best_option_perception = {}
  if config.goal:
    best_option_perception[goal_label] = goal_label
  best_option_perception.update({
      _get_class_name(observation): observation_label,
      _get_class_name(observation_summary): observation_summary_label,
      paranoia_label: paranoia_label,
      _get_class_name(relevant_memories): relevant_memories_label,
      _get_class_name(people_representation): person_representation_label,
      _get_class_name(options_perception): options_perception_label,
  })
  best_option_perception = (
      agent_components.question_of_recent_memories.BestOptionPerception(
          model=model,
          components=best_option_perception,
          clock_now=clock.now,
          pre_act_key=best_option_perception_label,
          logging_channel=measurements.get_channel(
              'BestOptionPerception'
          ).on_next,
      )
  )

  entity_components = (
      # Components that provide pre_act context.
      instructions,
      time_display,
      observation,
      observation_summary,
      relevant_memories,
      people_representation,
      options_perception,
      best_option_perception,
  )

  components_of_agent = {_get_class_name(component): component
                         for component in entity_components}
  components_of_agent[
      agent_components.memory_component.DEFAULT_MEMORY_COMPONENT_NAME] = (
          agent_components.memory_component.MemoryComponent(raw_memory))
  components_of_agent.update(additional_components)

  component_order = list(components_of_agent.keys())
  if overarching_goal is not None:
    components_of_agent[goal_label] = overarching_goal
    # Place goal after the instructions.
    component_order.insert(1, goal_label)

  components_of_agent[paranoia_label] = paranoia
  component_order.insert(
      component_order.index(_get_class_name(observation_summary)) + 1,
      paranoia_label)

  act_component = agent_components.concat_act_component.ConcatActComponent(
      model=model,
      clock=clock,
      component_order=component_order,
      logging_channel=measurements.get_channel('ActComponent').on_next,
  )

  agent = entity_agent_with_logging.EntityAgentWithLogging(
      agent_name=agent_name,
      act_component=act_component,
      context_components=components_of_agent,
      component_logging=measurements,
  )

  return agent
